package oktx.dubbo.config;

import com.alibaba.dubbo.common.extension.ExtensionLoader;
import com.alibaba.dubbo.config.annotation.Reference;
import com.alibaba.dubbo.rpc.cluster.LoadBalance;
import oktx.dubbo.loadbalance.DubboTxLoadBalanceWrapper;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.InstantiationAwareBeanPostProcessorAdapter;
import org.springframework.core.PriorityOrdered;
import org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;

import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import static org.springframework.core.BridgeMethodResolver.findBridgedMethod;
import static org.springframework.core.BridgeMethodResolver.isVisibilityBridgeMethodPair;
import static org.springframework.core.annotation.AnnotationUtils.findAnnotation;
import static org.springframework.core.annotation.AnnotationUtils.getAnnotation;

/**
 * Created by dave on 18-5-30 下午7:28.
 */
public class DubboBeanPostProcessor extends InstantiationAwareBeanPostProcessorAdapter implements PriorityOrdered {
    private Set<String> loadBalances = new HashSet<>();

    @Override
    public Object postProcessBeforeInstantiation(Class<?> beanClass, String beanName) throws BeansException {
        processMethodReferenceMetadata(beanClass);
        processFieldReferenceMetadata(beanClass);
        return null;
    }

    @Override
    public int getOrder() {
        return LOWEST_PRECEDENCE;
    }

    private void processFieldReferenceMetadata(final Class<?> beanClass) {
        ReflectionUtils.doWithFields(beanClass, field ->  {
            Reference reference = getAnnotation(field, Reference.class);
            if (reference != null) {
                updateReferenceParam(reference);
            }
        });

    }

    private void processMethodReferenceMetadata(final Class<?> beanClass) {
        ReflectionUtils.doWithMethods(beanClass, method ->  {
            Method bridgedMethod = findBridgedMethod(method);
            if (!isVisibilityBridgeMethodPair(method, bridgedMethod)) {
                return;
            }

            Reference reference = findAnnotation(bridgedMethod, Reference.class);
            if (reference != null && method.equals(ClassUtils.getMostSpecificMethod(method, beanClass))) {
                updateReferenceParam(reference);
            }
        });

    }

    private void updateReferenceParam(Reference ref) {
        String lb = ref.loadbalance();
        String key = "oktx";
        if (!lb.equals("")) {
            key = "oktx_" + lb;
            if (!loadBalances.contains(key)) {
                ExtensionLoader<LoadBalance> loader = ExtensionLoader.getExtensionLoader(LoadBalance.class);
                LoadBalance b = loader.getExtension(lb);
                loader.addExtension(key, DubboTxLoadBalanceWrapper.class);
                DubboTxLoadBalanceWrapper wrapper = (DubboTxLoadBalanceWrapper) loader.getExtension(key);
                wrapper.wrapInstanceLoadBalance(b);
                loadBalances.add(key);
            }
        }
        changeAnnotationValue(ref, "loadbalance", key);
    }

    @SuppressWarnings("unchecked")
    private static void changeAnnotationValue(Annotation annotation, String key, Object newValue) {
        Object handler = Proxy.getInvocationHandler(annotation);
        Field f;
        try {
            f = handler.getClass().getDeclaredField("memberValues");
        } catch (NoSuchFieldException | SecurityException e) {
            throw new IllegalStateException(e);
        }
        f.setAccessible(true);

        Map<String, Object> memberValues;
        try {
            memberValues = (Map<String, Object>) f.get(handler);
        } catch (IllegalArgumentException | IllegalAccessException e) {
            throw new IllegalStateException(e);
        }

        Object oldValue = memberValues.get(key);
        if (oldValue == null || oldValue.getClass() != newValue.getClass()) {
            throw new IllegalArgumentException();
        }
        memberValues.put(key,newValue);
    }
}
